import torchvision.models as models
"""
核心组件解析
"""
#1.torchvision.models
#加载预训练模型
resnet18=models.resnet18(pretrain=True)
alexnet=models.alexnet(pretrain=True)
vgg16=models.vgg16(pretrain=True)
#常用模型列表
#ResNet 通常用于图像分类 参数量：11M-60M top1准确率:69%-80%
#VGG 特征提取 参数量：138M top1准确率:71.3%
#MobileNet 移动端应用 3.4M top1准确率：70.6%
#EfficientNET 高效模型 5M-6M top1准确率：77%-84%

#2.torchvision.datasets
from torchvision import datasets
#加载CIFAR10数据集
train_data=datasets.CIFAR10(
    root='data',
    train=True,
    download=True,
    transform=transforms.ToTenso()
)

#加载MNIST数据集
test_data=datasets.MNIST(
    root='data',
    train=False,
    download=True
)

#支持的数据集类型
"""
graph TD
    A[torchvision.datasets] -- > B[分类数据集]
    A --> C[检测数据集]
    A --> D[分割数据集]
    B --> B1[CIFAR10/100]
    B --> B2[MNIST/FashingMNIST]
    B --> B3[ImageNet]
    C --> C1[COCO]
    C --> C2[VOC]
    D --> D1[Cityscapes]
"""

#3.torchvision.trainforms图像预处理和数据增强的核心工具
from torchvision import transforms
#定义图像变换管道
transform=transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485,0.456,0.406],
        std=[0.229,0.224,0.225]
    )
])
#常用的变换方法分类
#几何变换 RandomRotation,RandomResizedCrop 增加位置不变性
# 颜色变换 ColorJitter,Grayscale 增强颜色鲁棒性
# 模糊/噪声 GaussianBlur,RandomErasing 防止过拟合
# 组合变换 RandomApply,RandomChoice 灵活组合策略

"""
图像分类流程
""" 
#数据准备
import torch
from torchvision import datasets,transforms
from torch.utils.data import DataLoader

#定义数据变换
train_transform=transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(),
    transforms.ToTensor(),
    transforms.Normalize((0.5,),(0.5,))
])

#加载数据集
train_set=datasets.CIFAR10(
    root='./data',
    train=True,
    download=True,
    transform =train_transform
)

#创建数据加载器
train_loader=DataLoader(
    train_set,
    batch_size=32,
    shuffle=True
)

"""
模型训练
"""
import torch.nn as nn
import torch.optim as optim

#使用预训练模型
model=models.resnet18(pretrained=True)

#修改最后一层（适应CFIFAR10的10分类）
num_ftrs=model.fc.in_features
model.fc=nn.Linear(num_ftrs,10)

#定义损失函数和优化器
criterion=nn.CrossEntropyLoss()
optimizer=optim.SGD(model.parameters(),lr=0.0001,momentum=0.9)

#循环训练
for epoch in range(10):
    for images,labels in train_loader:
        outputs=model(images)
        loss=criterion(outputs,labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

"""
高级功能
"""
#1.自定义数据集
from torchvision.datasets import VisionDataset
class CustomDataset(VisionDataset):
    def __init__(self,root,transform=None):
        super().__init__(root,transform=transform)
        #实现 __getitem__ 和 __len__

    def __getitem__(self,index):
        #返回（image,target）元组
        pass
    def __len__(self):
        #返回数据集大小
        pass

#2.模型导出与部署
#导出为ONNX格式
dummy_input=torch.randn(1,3,224,224)
torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    input_name=['input'],
    output_name=['output']
)

"""
最佳使劲按建议
1.数据增强策略
①训练时随机变换增强数据
②验证/测试时候只使用确定性变换
"""
#2.迁移学习技巧
#冻结除最后一层外的所有参数
for param in model.parmeters():
    param.requires_grad=False
model.fc.requires_grad=True

#3.性能优化
#使用num_workers参数加速数据加载
#对大数据集考虑使用Dataset的子集

#4.常见错误
#忘记调用zero_grad()
#混淆了train()和eval()模式
#图像张量形状不符合模型要求(应为CxHxW)





